# -*-coding:utf-8-*-
# 生成训练验证集和测试集列表

import pdb
import glob
import os
import random
import math

# 这个函数是传入要比较的标签类别和查询原图是否有这个标签
def get_sample_value(txt_name, category_name):
    # kitti数据集的label位置  E:\DATASET\vocDataset\training\label_2
    label_path = 'E:/DATASET/vocDataset/training/label_2/'
    # 建立路径，用于记录哪个训练/测试用的数据集搭配哪个种类的文件名
    txt_path = label_path + txt_name + '.txt'
    try:
        with open(txt_path) as r_tdf:
            # 如果当前的这个种类在label里有则返回1，用于上层记录是否归属该类别
            if category_name in r_tdf.read():
                return ' 1'
            else:
                return '-1'
    except IOError as ioerr:
        print('File error:' + str(ioerr))


# glob文件搜索，找到符合这个路径的路径
# 也有博客有这个写法：txt_list_path = glob.glob('./Labels/*.txt')
txt_list_path = glob.glob('E:/DATASET/vocDataset/training/label_2/*.txt')
# 用于下文记录txt的所有文件名
txt_list = []

for item in txt_list_path:
    # 分离路径名和文件名
    temp1, temp2 = os.path.splitext(os.path.basename(item))
    # 将文件名加入到这个数组里
    txt_list.append(temp1)
# 对文件名进行排序
txt_list.sort()
# 控制台里显示所有的文件名
print(txt_list, end='\n\n')

# 有博客建议train:val:test=8:1:1，先尝试用一下
# 思路是将trainval分得9份，8份给train，1份给val。
# val用于训练时候验证集，也就是训练结束时得到的准确率。和测试集不一样

# random.sample多用于截取列表的指定长度的随机数，但是不会改变列表本身的排序
# math.floor是得到小于或者等于参数的最大整数
num_trainval = random.sample(txt_list, math.floor(len(txt_list) * 9 / 10.0))  # 可修改百分比
num_trainval.sort()
print(num_trainval, end='\n\n')

num_train = random.sample(num_trainval, math.floor(len(num_trainval) * 8 / 9.0))  # 可修改百分比
num_train.sort()
print(num_train, end='\n\n')

# 这里用set().difference()得到和刚才得到trian数据集不一样的部分，也就是9份里的剩下1份
num_val = list(set(num_trainval).difference(set(num_train)))
num_val.sort()
print(num_val, end='\n\n')

num_test = list(set(txt_list).difference(set(num_trainval)))
num_test.sort()
print(num_test, end='\n\n')

# pdb.set_trace()

# 自己创建VOC里imagesets的Main_path文件夹路径  E:\DATASET\ImageSets\Main
Main_path = 'E:/DATASET/ImageSets/Main/'
# 数据集种类，用于区分训练和测试用
train_test_name = ['trainval', 'train', 'val', 'test']
# 数据集里的各个类别
category_name = ['van', 'tram', 'car', 'pedestrian', 'truck']  # 修改类别

# 循环写trainvl train val test
for item_train_test_name in train_test_name:
    list_name = 'num_'
    # 这个num_加上数据集种类，可以得到上文8:1:1里的各个数据集名称
    list_name += item_train_test_name
    # 创建数据集种类的文件，记录这个种类都有哪些照片
    train_test_txt_name = Main_path + item_train_test_name + '.txt'
    try:
        # 写单个文件
        with open(train_test_txt_name, 'w') as w_tdf:
            # 一行一行写，eval是读取这个list的内容，也就是数据集种类里的照片名
            # 最终能得到四类：trainval.txt之类的四个
            for item in eval(list_name):
                w_tdf.write(item + '\n')
        # 循环写Car Pedestrian Cyclist
        for item_category_name in category_name:
            # 得到：路径+car+_+trainval+.txt的文件部分，
            # 也就是数据类别在当前数据集的分类数据判断，如00001 1为含有该类00002 -1为不含该类
            category_txt_name = Main_path + item_category_name + '_' + item_train_test_name + '.txt'
            with open(category_txt_name, 'w') as w_tdf:
                # 一行一行写，记录照片标号和该照片是否属于该类别
                for item in eval(list_name):
                    w_tdf.write(item + ' ' + get_sample_value(item, item_category_name) + '\n')
    except IOError as ioerr:
        print('File error:' + str(ioerr))
